#include "kpextOpImgF32C1ROI.h"

using namespace kpext;

OpImgF32C1ROI::OpImgF32C1ROI(kpextVkContext *ctx) :
    BaseCompute(ctx)
{
    if(!m_vkCtx){
        throw std::runtime_error(__FUNCTION__ " construct with null pointer");
    }

    m_cols = 1;
    m_rows = 1;
    m_roi_cols = 1;
    m_roi_rows = 1;
    m_roi_x = 0;
    m_roi_y = 0;
    m_inBuf = nullptr;
    m_outBuf = nullptr;

    m_postBarrier = true;

    //  初始化描述符集及布局
    m_descrPool = m_vkCtx->descriptorPool(
                {2},
                1);

    vk::DescriptorSetLayoutBinding binding(
                0,
                vk::DescriptorType::eStorageBuffer,
                1,
                vk::ShaderStageFlagBits::eCompute
                );
    std::vector<vk::DescriptorSetLayoutBinding> bindings;
    bindings.push_back(
                binding
                );

    binding.setBinding(1);
    bindings.push_back(binding);

    m_descrSetlayout = m_descrPool->createDescriptorSetLayout(
                bindings
                );
    m_descrSet = m_descrPool->recreateDescSet(*m_descrSetlayout,
                                              m_descrSet
                                              );

    // 初始化pipeline
    m_pipeline = m_vkCtx->computePipeline();
    m_pipeline->m_specConsts = std::make_shared<PipelineConstSZ>(3, 4);
    m_pipeline->m_pushConsts = std::make_shared<PipelineConstSZ>(6, 4);

    m_workGroupSize = ctx->getDefWorkGroupSize();

    m_pipeline->m_specConsts->set({m_workGroupSize[0],
                                   m_workGroupSize[1],
                                   1
                                  });
    m_pipeline->m_pushConsts->set({1, 1, 1, 1, 0, 0});

    m_pipeline->m_descrSetLayouts.push_back(m_descrSetlayout);

    if(!m_pipeline->rebuild(
                sizeof (spirv::ImgF32C1ROI), spirv::ImgF32C1ROI)
            ){
        SPDLOG_WARN("Rebuild kpext::ComputePipeline failed !");
    }

    //  工作组数量
    m_workCount = calc2DWorkGroup(1, 1,
                                  m_workGroupSize);
}

OpImgF32C1ROI::~OpImgF32C1ROI()
{

}

void OpImgF32C1ROI::rebuild(int cols, int rows,
                            int roi_cols, int roi_rows,
                            int roi_x, int roi_y,
                            std::shared_ptr<BaseBuffer> inbuf,
                            std::shared_ptr<BaseBuffer> outbuf)
{
    if(!inbuf || !outbuf || !m_vkCtx){
        throw std::runtime_error(__FUNCTION__ " input with null pointer!");
    }

    if(cols < 0) cols = 1;
    if(rows < 0) rows = 1;
    if(roi_cols < 0) roi_cols = 1;
    if(roi_rows < 0) roi_rows = 1;
    if(roi_x < 0) roi_x = 0;
    if(roi_y < 0) roi_y = 0;

    { // 检查ROI范围是否超出边界
        bool err = false;
        if(roi_x >= cols){
            err = true;
            roi_x = 0;
        }
        if(roi_y >= rows){
            err = true;
            roi_y = 0;
        }

        if(roi_x + roi_cols > cols){
            err = true;
            roi_cols = cols - roi_x;
        }
        if(roi_y + roi_rows > rows){
            err = true;
            roi_rows = rows - roi_y;
        }

        char const *  const errstr = { __FUNCTION__ " ROI is out of input boundary !"};
        if(err){
            SPDLOG_ERROR(errstr);
        }
    }

    m_cmdBarrier.prepare(
                {outbuf,},
                vk::PipelineStageFlagBits::eComputeShader,
                vk::PipelineStageFlagBits::eComputeShader,
                vk::AccessFlagBits::eShaderWrite,
                vk::AccessFlagBits::eShaderRead
                );

    STCAST_ASSIGN(m_cols, cols);
    STCAST_ASSIGN(m_rows, rows);
    STCAST_ASSIGN(m_roi_cols, roi_cols);
    STCAST_ASSIGN(m_roi_rows, roi_rows);
    STCAST_ASSIGN(m_roi_x, roi_x);
    STCAST_ASSIGN(m_roi_y, roi_y);

    if(m_inBuf != inbuf.get()
            || m_outBuf != outbuf.get()
            ){
        BaseCompute::rebuild({inbuf, outbuf});
        m_inBuf = inbuf.get();
        m_outBuf = outbuf.get();
    }

    //  设定输入输出尺寸
    m_pipeline->m_pushConsts->set(
                {m_cols, m_rows,
                 m_roi_cols, m_roi_rows,
                 m_roi_x, m_roi_y
                }
                );

    //  重新计算工作组大小
    m_workCount = calc2DWorkGroup(m_roi_cols, m_roi_rows,
                                  m_workGroupSize);
}
